from zhipuai import ZhipuAI
import json

MAX_RETRY = 5
API_KEY = "64b779db18b75b6ff548fa5c67ac1224.rcSgbAX6xnvkE4RM"


def LLM(query):
    client = ZhipuAI(api_key=API_KEY)
    response = client.chat.completions.create(
        model="glm-4",
        messages=[
            {"role": "user", "content": query},
        ],
        stream=False,
    )
    return response.choices[0].message.content


def get_json_response(query, max_retries=MAX_RETRY):
    response = LLM(query)
    for retry in range(max_retries):
        try:
            response = LLM(query)
            # print(response)
            response = json.loads(response)
            return response
        except Exception as e:
            print(f"解析失败，正在进行第{retry + 1}次重试，错误信息：{e}")
            if retry == max_retries - 1:
                print(f"解析失败，已达到最大重试次数{max_retries}。")
                return None


def refine_answer(question, answer):
    prompt = f"""
    问题：{question}
    信息：{answer}
    请整合答案，直接给出简洁、完整且清晰的回答。回答格式忠于提问方式，在答案前面带上问题，比如请问(2019）鲁0991执保239号的上诉人和被诉人分别是？则回答为：(2019）鲁0991执保239号的上诉人是XXXX，被诉人是XXXX。
    不要回答问题之外的内容，也不要回复"回答"这两个字，也无需进一步解释。
    回答：
    """
    final_answer = LLM(query=prompt)
    return final_answer


def refine_question(question):
    prompt = f"""根据例子，完善问题
例子：劲拓股份拥有哪些子公司
完善：关联上市公司全称是劲拓股份的公司拥有哪些子公司

例子：华仁药业股份有限公司控股的子公司中，持股比例超过50%的有几家？
完善：关联上市公司全称是华仁药业股份有限公司控股的子公司中，持股比例超过50%的有几家？

例子：北京华清瑞达科技有限公司、博晖生物制药（内蒙古）有限公司、浙江迪安健检医疗管理有限公司分别属于哪几家公司的子公司
完善：北京华清瑞达科技有限公司、博晖生物制药（内蒙古）有限公司、浙江迪安健检医疗管理有限公司分别属于哪几家公司的子公司

例子：Shanghai Construction Group Co., Ltd.拥有哪些子公司？
完善：关联上市公司全称是Shanghai Construction Group Co., Ltd.拥有哪些子公司？
----
问题：{question}
完善：(不要回答问题)"""
    return LLM(prompt)


def get_court_name(question):
    prompt = f"""根据问题，提取法院名称
例子：上诉于江西省赣州市中级人民法院 
法院名称：江西省赣州市中级人民法院

----
问题：{question}
法院名称：(不要回答问题)"""
    return LLM(prompt)


def get_province_city_area(question):
    prompt = f"""根据问题，将地址拆分为：省份、城市、区县
例子：江苏省连云港市高新开发区，拆分为："省份":"江苏省","城市":"连云港市","区县":"高新开发区"
不要多余的答案，只需要按格式回答即可

----
问题：{question}
"""
    result = LLM(prompt)
    result = '{' + result + "}"
    dict_result = json.loads(result)
    return dict_result

def completion_case_num(question):
    prompt = f"""根据不符合格式的案号，修正补全案号格式
    例子：2019苏0214民初6847号 
    案号：(2019)苏0214民初6847号
    不要多余的答案，只需要按格式回答即可

    ----
    问题：{question}"""
    return LLM(prompt)

def completion_court_name(court_name):
    prompt = f"""根据不符合格式的法院名称，修正法院名称
    例子：福建漳州中级人民法院 
    法院名称：福建省漳州市中级人民法院
    不要多余的答案，只需要按格式回答即可

    ----
    问题：{court_name}"""
    return LLM(prompt)


if __name__ == "__main__":
    print(completion_court_name("福建漳州中级人民法院"))

